
// MNIST 画板交互逻辑
const canvas = document.getElementById('mnist-canvas');
const ctx = canvas.getContext('2d');
let isDrawing = false;

// 初始化画板
function initCanvas() {
  ctx.fillStyle = 'black';
  ctx.fillRect(0, 0, canvas.width, canvas.height);
  ctx.lineWidth = 15;
  ctx.lineCap = 'round';
  ctx.strokeStyle = 'white';
}

// 绘图事件处理
canvas.addEventListener('mousedown', startDrawing);
canvas.addEventListener('mousemove', draw);
canvas.addEventListener('mouseup', stopDrawing);
canvas.addEventListener('mouseout', stopDrawing);

function startDrawing(e) {
  isDrawing = true;
  draw(e);
}

function draw(e) {
  if (!isDrawing) return;
  ctx.lineTo(e.offsetX, e.offsetY);
  ctx.stroke();
  ctx.beginPath();
  ctx.moveTo(e.offsetX, e.offsetY);
}

function stopDrawing() {
  isDrawing = false;
  ctx.beginPath();
}

// 清除画板
document.getElementById('clear-btn').addEventListener('click', () => {
  initCanvas();
  document.getElementById('mnist-result').textContent = '';
});

// 加载MNIST模型并进行预测
async function loadMNISTModel() {
  const model = await tf.loadLayersModel('https://storage.googleapis.com/tfjs-models/tfjs/mnist_cnn/model.json');
  return model;
}

document.getElementById('predict-btn').addEventListener('click', async () => {
  const model = await loadMNISTModel();
  const tensor = preprocessCanvas();
  const predictions = model.predict(tensor).dataSync();
  const predictedNum = predictions.indexOf(Math.max(...predictions));
  document.getElementById('mnist-result').textContent = `预测数字: ${predictedNum}`;
});

function preprocessCanvas() {
  const tempCanvas = document.createElement('canvas');
  const tempCtx = tempCanvas.getContext('2d');
  tempCanvas.width = 28;
  tempCanvas.height = 28;
  tempCtx.drawImage(canvas, 0, 0, 28, 28);
  const imgData = tempCtx.getImageData(0, 0, 28, 28);
  return tf.tidy(() => {
    const tensor = tf.browser.fromPixels(imgData, 1)
      .resizeNearestNeighbor([28, 28])
      .toFloat()
      .div(255.0)
      .expandDims();
    return tensor;
  });
}

// 线性回归演示
const regressionCtx = document.getElementById('regression-plot').getContext('2d');
const regressionChart = new Chart(regressionCtx, {
  type: 'scatter',
  data: {
    datasets: [{
      label: '训练数据',
      borderColor: 'rgb(75, 192, 192)',
      backgroundColor: 'rgba(75, 192, 192, 0.5)',
      pointRadius: 8,
      data: []
    }, {
      label: '预测线',
      borderColor: 'rgb(255, 99, 132)',
      backgroundColor: 'rgba(255, 99, 132, 0.2)',
      pointRadius: 0,
      data: [],
      type: 'line',
      fill: false,
      borderWidth: 2
    }]
  },
  options: {
    responsive: true,
    scales: {
      x: { title: { display: true, text: '输入值' } },
      y: { title: { display: true, text: '输出值' } }
    }
  }
});

// 生成训练数据
function generateData() {
  const xs = tf.tensor2d([1, 2, 3, 4], [4, 1]);
  const ys = tf.tensor2d([1.1, 2.9, 5.2, 7.1], [4, 1]);
  
  // 更新图表
  const xValues = xs.dataSync();
  const yValues = ys.dataSync();
  const points = xValues.map((x, i) => ({x, y: yValues[i]}));
  regressionChart.data.datasets[0].data = points;
  regressionChart.update();
  
  return {xs, ys};
}

// 创建线性回归模型
function createModel() {
  const model = tf.sequential();
  model.add(tf.layers.dense({units: 1, inputShape: [1]}));
  model.compile({loss: 'meanSquaredError', optimizer: 'sgd'});
  return model;
}

// 训练模型
document.getElementById('train-btn').addEventListener('click', async () => {
  const model = createModel();
  const {xs, ys} = generateData();
  
  await model.fit(xs, ys, {
    epochs: 100,
    callbacks: {
      onEpochEnd: (epoch, logs) => {
        if (epoch % 10 === 0) {
          updateRegressionLine(model);
          document.getElementById('regression-result').innerHTML = 
            `训练中... 第${epoch+1}轮 损失值: ${logs.loss.toFixed(4)}`;
        }
      }
    }
  });
  
  document.getElementById('regression-result').innerHTML = 
    '训练完成! 尝试预测X=5时: ' + model.predict(tf.tensor2d([5], [1, 1])).dataSync()[0].toFixed(2);
});

// 更新回归线
function updateRegressionLine(model) {
  const xValues = [-1, 0, 1, 2, 3, 4, 5];
  const predictedYs = model.predict(tf.tensor2d(xValues, [xValues.length, 1])).dataSync();
  const linePoints = xValues.map((x, i) => ({x, y: predictedYs[i]}));
  regressionChart.data.datasets[1].data = linePoints;
  regressionChart.update();
}

// 初始化
initCanvas();
generateData();
